package com.boge.ai.utils;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import okhttp3.*;
import org.apache.commons.math3.linear.RealVector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Milvus 向量数据库的公共方法
 */
public class MilvusUtils {

/*    private final static String BASE_HOST = "localhost";
    private final static int BASE_PORT = 19530;*/
    private  String collectionName = "llm_rag";
    private MilvusServiceClient client ;

    public MilvusUtils(String host, int port) {
        this(host,port,null);
    }

    public MilvusUtils(String host, int port, String collectionName) {
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost(host)
                .withPort(port)
                .build();
        client = new MilvusServiceClient(connectParam);
        if(collectionName != null) {
            this.collectionName = collectionName;
        }

    }

    /**
     * 创建集合
     * @throws Exception
     */
    public  void createCollection() throws Exception {

        List<FieldType> fieldTypes = Arrays.asList(
                FieldType.newBuilder()
                        .withName("id")
                        .withDataType(DataType.Int64)
                        .withPrimaryKey(true)
                        .withAutoID(false)
                        .build(),
                FieldType.newBuilder()
                        .withName("vector")
                        .withDataType(DataType.FloatVector)
                        .withDimension(128) // 设置向量维度
                        .build());

        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("Test collection")
                .withShardsNum(2)
                .withFieldTypes(fieldTypes)
                .build();
        client.createCollection(createCollectionReq);
        // 同时给向量创建对应的索引
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName("vector") // 向量字段名
                .withIndexType(IndexType.IVF_FLAT) // 使用IVF_FLAT索引类型
                .withMetricType(MetricType.L2) // 指定度量类型，如L2距离
                .withExtraParam("{\"nlist\":128}") // 根据索引类型提供额外参数，比如nlist
                .build();

        client.createIndex(createIndexParam);
    }

    /**
     * 插入数据到向量数据库
     * @throws Exception
     */
    public  void insertVectoryData(List<Long> idList,List<List<Float>> vectors) throws Exception {

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("id", idList));
        fields.add(new InsertParam.Field("vector", vectors));


        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fields)
                .build();
        client.insert(insertParam);
    }

    /**
     * 插入数据到向量数据库
     * @throws Exception
     */
    public  void insertVectoryLLMData(List<String> idList,List<RealVector> vectors) throws Exception {

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("id", idList));
        fields.add(new InsertParam.Field("vector", vectors));


        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fields)
                .build();
        client.insert(insertParam);
    }

    /**
     * 根据向量检索信息
     * @param searchVectors
     * @return
     * @throws Exception
     */
    public  SearchResultsWrapper search(List<List<Float>> searchVectors) throws Exception {

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withMetricType(MetricType.L2)// 使用 L2 距离作为相似度度量
                .withTopK(3) // 返回最接近的前3个结果
                .withVectors(searchVectors)
                .withVectorFieldName("vector") // 向量字段名
                .withOutFields(Arrays.asList("id")) // 需要返回的字段
                .build();
        SearchResults data = client.search(searchParam).getData();
        if(data != null) {
            SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());
            resultsWrapper.getRowRecords().forEach(result -> {
                System.out.println("Search result: " + result);
            });
            return resultsWrapper;
        }
        return null;
    }

    public  SearchResultsWrapper searchLLM(List<RealVector> searchVectors) throws Exception {

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withMetricType(MetricType.L2)// 使用 L2 距离作为相似度度量
                .withTopK(3) // 返回最接近的前3个结果
                .withVectors(searchVectors)
                .withVectorFieldName("vector") // 向量字段名
                .withOutFields(Arrays.asList("id")) // 需要返回的字段
                .build();
        SearchResults data = client.search(searchParam).getData();
        if(data != null) {
            SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());
            resultsWrapper.getRowRecords().forEach(result -> {
                System.out.println("Search result: " + result);
            });
            return resultsWrapper;
        }
        return null;
    }

    public long getTextEmbedding(String msg) throws Exception {
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(20, TimeUnit.SECONDS)
                .readTimeout(20, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, String.format("""
                {
                  "input": "%s",
                  "model": "text-embedding-3-small"
                }
                """,msg));
        Request request = new Request.Builder()
                .url("https://ai-yyds.com/v1/embeddings")
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .addHeader("Authorization", "Bearer sk-UrwVZPGfWFwqXGLRE51c1b4dBeD849CcB71dF8E5579a4a09")
                .build();
        Response response = client.newCall(request).execute();
        ResponseBody responseBody = response.body();
        String jsonString = responseBody.string();
        return Long.parseLong(jsonString);
    }

}
